#include <chrono>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <CL/cl.h>
#include <string.h>
#include <sys/time.h>

#define CHECK(err, expr) \
    do { \
        cl_int errCode = (expr); \
        if (errCode != err) { \
            fprintf(stderr, "Error: %d at %s:%d\n", errCode, __FILE__, __LINE__); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)


#define ROWS_MAX 5
class opencl_accelerate
{
private:
    // 定义矩阵尺寸
    size_t rowsA = ROWS_MAX, colsA = 32;
    size_t rowsB = colsA, colsB = 160*160;

    float* hostA;
    float* hostB;
    uint8_t* hostResult;
    cl_context context;
    cl_command_queue queue;
    cl_mem bufferA;
    cl_mem bufferB;
    cl_mem bufferResult;
    cl_program program;
    cl_kernel kernel;

    cl_int err;

    // OpenCL平台初始化
    cl_platform_id platform = NULL;
    cl_device_id device = NULL;

    // 内核源代码
    const char* kernel_src = 
    R"(
        __kernel void matrixMultiplication(__global float* A, __global float* B, __global uchar* C,const int widthA, const int heightA, const int widthB){
            int Row = get_global_id(0);
            int Col = get_global_id(1);
            if (Row < heightA && Col < widthB) {
                float Cvalue = 0.0f;
                for (int k = 0; k < widthA; ++k) {
                    Cvalue += A[Row * widthA + k] * B[k * widthB + Col];
                }
                if (Cvalue > 0){
                    C[Row * widthB + Col] = 4;
                }else{
                    C[Row * widthB + Col] = 0;
                }
            }
        }
    )";
public:
    void set_buffer()
    {
        // 创建缓冲区对象
        bufferA = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, rowsA * colsA * sizeof(float), hostA, &err);
        CHECK(CL_SUCCESS, err);
        
        bufferB = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, rowsB * colsB * sizeof(float), hostB, &err);
        CHECK(CL_SUCCESS, err);

        // 设置内核参数
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 0, sizeof(cl_mem), &bufferA));
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 1, sizeof(cl_mem), &bufferB));
    }
    opencl_accelerate()
    {
        // size_t rowsA = 4, colsA = 4;
        // size_t rowsB = colsA, colsB = 4;
        cl_uint num_platforms;
        // 初始化主机端矩阵A和B
        hostA = (float*)malloc(rowsA * colsA * sizeof(float));
        hostB = (float*)malloc(rowsB * colsB * sizeof(float));
        hostResult = (uint8_t*)malloc(rowsA * colsB * sizeof(uint8_t)); // 这里应该是一个参考结果

        // 随机初始化矩阵A和B
        for (size_t i = 0; i < rowsA * colsA; i++) hostA[i] = (float)i;
        for (size_t i = 0; i < rowsB * colsB; i++) hostB[i] = (float)i;
        // memset(hostB, 0, rowsB * colsB * sizeof(float));
        // memset(hostA, 0, rowsA * colsA * sizeof(float));
        

        CHECK(CL_SUCCESS, clGetPlatformIDs(1, &platform, &num_platforms));
        printf("num_platforms=%d\n", num_platforms);
        CHECK(CL_SUCCESS, clGetDeviceIDs(platform, CL_DEVICE_TYPE_DEFAULT, 1, &device, NULL));
        print_device_info(device);

        context = clCreateContext(NULL, 1, &device, NULL, NULL, &err);
        CHECK(CL_SUCCESS, err);
        
        queue = clCreateCommandQueueWithProperties(context, device, 0, &err);
        CHECK(CL_SUCCESS, err);

        // 创建缓冲区对象
        bufferA = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, rowsA * colsA * sizeof(float), hostA, &err);
        CHECK(CL_SUCCESS, err);
        
        bufferB = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, rowsB * colsB * sizeof(float), hostB, &err);
        CHECK(CL_SUCCESS, err);
        
        bufferResult = clCreateBuffer(context, CL_MEM_WRITE_ONLY, rowsA * colsB * sizeof(uint8_t), NULL, &err);
        CHECK(CL_SUCCESS, err);


        // 构建程序
        program = buildProgram(context, device,kernel_src);
        
        // 创建内核
        kernel = clCreateKernel(program, "matrixMultiplication", &err);
        CHECK(CL_SUCCESS, err);

        // 设置内核参数
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 0, sizeof(cl_mem), &bufferA));
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 1, sizeof(cl_mem), &bufferB));
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 2, sizeof(cl_mem), &bufferResult));
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 3, sizeof(int), &colsA));
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 4, sizeof(int), &rowsA));
        CHECK(CL_SUCCESS, clSetKernelArg(kernel, 5, sizeof(int), &colsB));
        std::cout << "kernel set args success" << std::endl;
    
        size_t global_work_size[2] = {rowsA, colsB};
        CHECK(CL_SUCCESS, clEnqueueNDRangeKernel(queue, kernel, 2, NULL, global_work_size, NULL, 0, NULL, NULL));
        CHECK(CL_SUCCESS, clEnqueueReadBuffer(queue, bufferResult, CL_TRUE, 0, rowsA * colsB * sizeof(uint8_t), hostResult, 0, NULL, NULL));
        std::cout << "Dump hostResult:" << std::endl;
        for (int i = 0; i < rowsA; i++)
        {
            // for (int j = 0; j < colsB; j++)
            // {
            //     printf("%.1f ", hostResult[j]);
            // }

            printf("%d ", hostResult[i*colsB]);
            printf("\n");
        }
    }
    ~opencl_accelerate()
    {
        // 清理资源
        CHECK(CL_SUCCESS, clReleaseMemObject(bufferA));
        CHECK(CL_SUCCESS, clReleaseMemObject(bufferB));
        CHECK(CL_SUCCESS, clReleaseMemObject(bufferResult));
        CHECK(CL_SUCCESS, clReleaseKernel(kernel));
        CHECK(CL_SUCCESS, clReleaseProgram(program));
        CHECK(CL_SUCCESS, clReleaseCommandQueue(queue));
        CHECK(CL_SUCCESS, clReleaseContext(context));
        free(hostA);
        free(hostB);
        free(hostResult);
    }
    cl_program buildProgram(cl_context context, cl_device_id device, const char* kernel_src) {
        cl_int err;
        cl_program program = clCreateProgramWithSource(context, 1, &kernel_src, NULL, &err);
        CHECK(CL_SUCCESS, err);
        CHECK(CL_SUCCESS, clBuildProgram(program, 1, &device, NULL, NULL, NULL));
        return program;
    }
    void print_device_info(cl_device_id device) {
       cl_int err;
       size_t size;

       // 获取设备名称长度
       err = clGetDeviceInfo(device, CL_DEVICE_NAME, 0, nullptr, &size);
       if (err != CL_SUCCESS) {
           std::cerr << "Error calling clGetDeviceInfo(CL_DEVICE_NAME): " << err << std::endl;
           return;
       }

       // 分配内存存储设备名称
       char* name = new char[size];

       // 获取设备名称
       err = clGetDeviceInfo(device, CL_DEVICE_NAME, size, name, nullptr);
       if (err != CL_SUCCESS) {
           std::cerr << "Error calling clGetDeviceInfo(CL_DEVICE_NAME): " << err << std::endl;
           delete[] name;
           return;
       }

       // 打印设备名称
       std::cout << "Device Name: " << name << std::endl;

       // 获取并打印其他设备信息
       clGetDeviceInfo(device, CL_DEVICE_VENDOR, size, name, nullptr);
       std::cout << "Vendor: " << name << std::endl;

       clGetDeviceInfo(device, CL_DEVICE_VERSION, size, name, nullptr);
       std::cout << "Version: " << name << std::endl;

       clGetDeviceInfo(device, CL_DRIVER_VERSION, size, name, nullptr);
       std::cout << "Driver Version: " << name << std::endl;

       cl_device_type type;
        clGetDeviceInfo(device, CL_DEVICE_TYPE, sizeof(type), &type, nullptr);
        std::cout << "Type: ";
        if (type & CL_DEVICE_TYPE_DEFAULT) std::cout << "Default ";
        if (type & CL_DEVICE_TYPE_CPU) std::cout << "CPU ";
        if (type & CL_DEVICE_TYPE_GPU) std::cout << "GPU ";
        if (type & CL_DEVICE_TYPE_ACCELERATOR) std::cout << "Accelerator ";
        std::cout << std::endl;

        cl_ulong max_compute_units;
        clGetDeviceInfo(device, CL_DEVICE_MAX_COMPUTE_UNITS, sizeof(max_compute_units), &max_compute_units, nullptr);
        std::cout << "Max Compute Units: " << max_compute_units << std::endl;

        cl_ulong max_work_item_sizes[3];
        clGetDeviceInfo(device, CL_DEVICE_MAX_WORK_ITEM_SIZES, sizeof(max_work_item_sizes), &max_work_item_sizes, nullptr);
        std::cout << "Max Work Item Sizes: "
                  << max_work_item_sizes[0] << ", "
                  << max_work_item_sizes[1] << ", "
                  << max_work_item_sizes[2] << std::endl;

        // 释放分配的内存
        delete[] name;
    }
    void calculat(float* A, float* B,uint8_t* C, int ROWS_A, int COLS_A, int COLS_B)
    {
        // auto start = std::chrono::high_resolution_clock::now();    // compute the mask through Matmul
        CHECK(CL_SUCCESS, clEnqueueWriteBuffer(queue, bufferA, CL_TRUE, 0, rowsA * colsA * sizeof(float), A, 0, NULL, NULL));
        CHECK(CL_SUCCESS, clEnqueueWriteBuffer(queue, bufferB, CL_TRUE, 0, rowsB * colsB * sizeof(float), B, 0, NULL, NULL));

        // // 复制矩阵A和B到主机端
        // printf("hostA size: %d, hostB size: %d\n", ROWS_A * COLS_A*sizeof(float), COLS_A * COLS_B*sizeof(float));
        // // memcpy(hostA, A, ROWS_A * COLS_A * sizeof(float));
        // // memcpy(hostB, B, COLS_A * COLS_B * sizeof(float));
        // // set_buffer();
        // printf("hostA:\n");
        // for (int i = 0; i < ROWS_A; i++)
        // {
        //     for (int k = 0; k < COLS_A; k++)
        //     {
        //         printf("%.1f ", hostA[k]);
        //     }
        //     printf("\n");
        // }
        // // printf("hostB:\n");
        // // for (int i = 0; i < COLS_A; i++)
        // // {
        // //     for (int k = 0; k < COLS_B; k++)
        // //     {
        // //         printf("%.1f ", hostB[k]);
        // //     }
        // //     printf("\n");
        // // }
        // 执行内核
        size_t global_work_size[2] = {rowsA, colsB};
        CHECK(CL_SUCCESS, clEnqueueNDRangeKernel(queue, kernel, 2, NULL, global_work_size, NULL, 0, NULL, NULL));
        // auto end = std::chrono::high_resolution_clock::now();
        // std::cout << "clEnqueueNDRangeKernel => duration:" << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() << "us" << std::endl;

        // 读取结果
        CHECK(CL_SUCCESS, clEnqueueReadBuffer(queue, bufferResult, CL_TRUE, 0, rowsA * colsB * sizeof(uint8_t), hostResult, 0, NULL, NULL));
        // for (int i = 0; i < rowsA; i++)
        // {
        //     // for (int j = 0; j < colsB; j++)
        //     // {
        //     //     printf("%.1f ", hostResult[j]);
        //     // }
        //     printf("%d ", hostResult[i*colsB]);
        //     printf("\n");
        // }
        
        memcpy(C, hostResult, ROWS_A * COLS_B * sizeof(uint8_t));

    }
};
